#!/usr/bin/env python
from __future__     import division
from collections    import namedtuple
from euclid         import Ray3, Vector3, LineSegment3
from geometry       import *
from math           import sqrt

MAX_REFLECTIONS = 2 
AMBIENT_COEF = 0.2
DIFFUSE_COEF = 1 - AMBIENT_COEF

class Scene:
    def __init__(self):
        self.solids = set()
    
    def append(self, solid):
        self.solids.add(solid)

    def intersection(self, ray, remove=None):
        solids = self.solids
        if remove is not None:
            solids = solids - remove

        intersection_points = []
        for obj in solids:
            inters = obj.intersect(ray)
            for inter in inters:
                #if inter[0] is not None and inter[0] != ray.p1:
                if inter[0] != ray.p1:
                    intersection_points.append(inter + (obj,))

        if not intersection_points:
            return (None, None, None)
        return min(intersection_points, key=lambda p: (ray.p1-p[0]).magnitude_squared())

    def ray_color(self, ray, n = MAX_REFLECTIONS, before = None):
        intersection, normal, solid = self.intersection(ray)
        if intersection is None:
            return BLACK
        if normal.dot(ray.v) > 0:
            # ray is exiting surface
            print "Ray out solid={0}, before={1}".format(solid, before)
            refracted_ray = self.refracted_ray(ray, intersection, -normal, solid, None)
            if refracted_ray is None:
                print "The ray couldn't get out of the solid o.O"
                return BLACK
            return self.ray_color(refracted_ray, n, None)
        color = self.color_at_point(intersection, normal, solid)
        if n > 0:
            if solid.reflection > 0:
                reflection_v = (ray.v - 2*ray.v.dot(normal)*normal).normalize()
                reflection_ray = Ray3(intersection, reflection_v)
                new_color = self.ray_color(reflection_ray, n-1, before)
                color += new_color*solid.reflection
            if solid.refraction > 0:
                print "Ray in"
                refraction_ray = self.refracted_ray(ray, intersection, normal, before, solid)
                if refraction_ray is not None:
                    new_color = self.ray_color(refraction_ray, n-1, solid)
                    print "ref_ray.dot(normal)={0}".format(refraction_ray.v.normalize().dot(normal))
                    color += new_color*solid.refraction

        return color

    @staticmethod
    def refracted_ray(ray, intersection, normal, before, after):
        v = ray.v.normalize()
        print "Before: {0}".format(v)
        normal = normal.normalize()
        c1 = normal.dot(v)
        n1 = 1
        if before is not None:
            n1 = before.ref_coefficient
        n2 = 1
        if after is not None:
            n2 = after.ref_coefficient
        assert n2>0, "Trying to exit solid {0}".format(after)
        n = n1/n2
        inside_term = 1-n**2*(1-c1**2)
        if inside_term <= 0:
            return None
        c2 = sqrt(inside_term)
        refraction_v = (n*v) - (n*c1+c2)*normal
        refraction_ray = Ray3(intersection, refraction_v)
        print "After: {0}".format(refraction_ray)
        return refraction_ray
    
    def color_at_point(self, intersection, normal, solid):
        color = solid.color*AMBIENT_COEF

        shadow_ray = LineSegment3(intersection, self.light)
        blocking_shadow_ray, _, _ = self.intersection(shadow_ray, set([solid]))
        if blocking_shadow_ray is None:
            light_vector = (self.light - intersection).normalize()
            shade = light_vector.dot(normal)
            shade = max(shade, 0)
            color += solid.color*DIFFUSE_COEF*shade

        return color


class Solid:
    def __init__(self, geometric_form, color, reflection=0, refraction=0, ref_coefficient=0):
        self.geometric_form = geometric_form
        self.color = color
        self.reflection = reflection
        self.refraction = refraction
        self.ref_coefficient = ref_coefficient

    def intersect(self, ray):
        inter = self.geometric_form.intersect(ray)
        assert isinstance(inter, list), "intersect ({0}) did not return list (returned {1}).".format(self.geometric_form, inter)
        return inter 

    def __str__(self):
        return "Solid(" + str(self.geometric_form) + ")"


class Color(namedtuple("Color", "r g b")):
    def __add__(self, other):
        return Color(
                min(self.r + other.r, 255),
                min(self.g + other.g, 255),
                min(self.b + other.b, 255)
                )

    def __mul__(self, scalar):
        return Color(
                min(self.r * scalar, 255),
                min(self.g * scalar, 255),
                min(self.b * scalar, 255)
                )

BLACK = Color(0, 0, 0)

